﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Xml;

using System.Diagnostics.Contracts;
using SimpleIR;

namespace PInfoModel
{
	public sealed class PIModel
	{
		/// <summary>
		/// True if the state represented by this model is infeaseable (can never occour).
		/// </summary>
		public readonly bool InfeasableState;

		/// <summary>
		/// A map from locals and statics to the <c>MNode</c> that represents the value of the name. 
		/// </summary>
		public readonly Dictionary<IRRootIdentifier, int> NameToNodeMap;

		/// <summary>
		/// A map from external references (and the parameter targets at call) to nodes in this..
		/// </summary>
		public readonly Dictionary<PIExternalRef, int[]> ExternalRefs;

		/// <summary>
		/// A list of <c>SNode</c> objects representing the eval stack, grows from 0 up.
		/// </summary>
		public readonly List<PISNode> EvalStack;

		/// <summary>
		/// All the nodes in the graph.
		/// </summary>
		public readonly List<PIMNode> Nodes;

		[ContractInvariantMethod]
		private void PModelInvariant()
		{
			Contract.Invariant(this.NameToNodeMap != null);
			Contract.Invariant(this.ExternalRefs != null);
			Contract.Invariant(this.EvalStack != null);

			Contract.Invariant(this.Nodes != null);
		}

		public PIModel()
		{
			this.InfeasableState = false;
			this.NameToNodeMap = new Dictionary<IRRootIdentifier, int>();
			this.ExternalRefs = new Dictionary<PIExternalRef, int[]>();
			this.EvalStack = new List<PISNode>();

			this.Nodes = new List<PIMNode>();
		}

		public PIModel(bool isinf, Dictionary<IRRootIdentifier, int> rootmap, Dictionary<PIExternalRef, int[]> extmap, List<PISNode> stack, List<PIMNode> nodes)
		{
			this.InfeasableState = isinf;
			this.NameToNodeMap = rootmap;
			this.ExternalRefs = extmap;
			this.EvalStack = stack;

			this.Nodes = nodes;
		}

		/// <summary>
		/// At initialization add the home location info for a given root.
		/// </summary>
		public void AddRootIdentifierHome(IRRootIdentifier rid, int target)
		{
			this.NameToNodeMap.Add(rid, target);
		}

		/// <summary>
		/// As initialization add the external ref info for a single extref.
		/// </summary>
		public void AddExternalRefTrgets(PIExternalRef extr, int[] targets)
		{
			this.ExternalRefs.Add(extr, targets);
		}

		internal void Serialize(PInfoSerializer sr)
		{
			sr.WriteStartTag("PModel");

			sr.WriteBool("Inf", this.InfeasableState);
			
			sr.WriteCollectionStartTag("NameMap", this.NameToNodeMap.Count);
			foreach(KeyValuePair<IRRootIdentifier, int> nnmp in this.NameToNodeMap)
			{
				sr.WriteStartTag("NME");
				sr.WriteIRRoot("Name", nnmp.Key);
				sr.WriteInt("Trgt", nnmp.Value);
				sr.WriteEndTag("NME");
			}
			sr.WriteCollectionEndTag("NameMap");

			sr.WriteCollectionStartTag("ExtRefs", this.ExternalRefs.Count);
			foreach(KeyValuePair<PIExternalRef, int[]> erp in this.ExternalRefs)
			{
				sr.WriteStartTag("ERE");
				erp.Key.Serialize(sr);
				sr.WriteCollectionStartTag("Trgts", erp.Value.Length);
				for(int i = 0; i < erp.Value.Length; ++i)
				{
					sr.WriteStartTag("Trgt");
					sr.WriteInt("Addr", erp.Value[i]);
					sr.WriteEndTag("Trgt");
				}
				sr.WriteCollectionEndTag("Trgts");
				sr.WriteEndTag("ERE");
			}
			sr.WriteCollectionEndTag("ExtRefs");

			sr.WriteCollectionStartTag("EvalStack", this.EvalStack.Count);
			for(int i = 0; i < this.EvalStack.Count; ++i)
				this.EvalStack[i].Serialize(sr);
			sr.WriteEndTag("EvalStack");

			sr.WriteCollectionStartTag("Nodes", this.Nodes.Count);
			for(int i = 0; i < this.Nodes.Count; ++i)
				this.Nodes[i].Serialize(sr);
			sr.WriteCollectionEndTag("Nodes");

			sr.WriteEndTag("PModel");
		}

		internal static PIModel Deserialize(PInfoDeSerializer dsr)
		{
			dsr.ReadStartTag("PModel");

			bool inf = dsr.ReadBool("Inf");
			
			Dictionary<IRRootIdentifier, int> namemap = new Dictionary<IRRootIdentifier,int>();
			int nmcount = dsr.ReadCollectionStartTag("NameMap");
			for(int i = 0; i < nmcount; ++i)
			{
				dsr.ReadStartTag("NME");
				IRRootIdentifier rrid = dsr.ReadIRRoot("Name");
				int trgt = dsr.ReadInt("Trgt");
				namemap.Add(rrid, trgt);
				dsr.ReadEndTag("NME");
			}
			dsr.ReadCollectionEndTag("NameMap");

			Dictionary<PIExternalRef, int[]> extmap = new Dictionary<PIExternalRef, int[]>();
			int extcount = dsr.ReadCollectionStartTag("ExtRefs");
			for(int i = 0; i < extcount; ++i)
			{
				dsr.ReadStartTag("ERE");
				PIExternalRef eref = PIExternalRef.Deserialize(dsr);

				int trgtcount = dsr.ReadCollectionStartTag("Trgts");
				int[] ta = new int[trgtcount];
				for(int j = 0; j < trgtcount; ++j)
				{
					dsr.ReadStartTag("Trgt");
					ta[j] = dsr.ReadInt("Addr");
					dsr.ReadEndTag("Trgt");
				}
				dsr.ReadCollectionEndTag("Trgts");

				dsr.ReadEndTag("ERE");
			}
			dsr.ReadCollectionEndTag("ExtRefs");

			List<PISNode> evalstack = new List<PISNode>();
			int stackcount = dsr.ReadCollectionStartTag("EvalStack");
			for(int i = 0; i < stackcount; ++i)
				evalstack.Add(PISNode.Deserialize(dsr));
			dsr.ReadCollectionEndTag("EvalStack");

			List<PIMNode> nodes = new List<PIMNode>();
			int nodecount = dsr.ReadCollectionStartTag("Nodes");
			for(int i = 0; i < nodecount; ++i)
				nodes.Add(PIMNode.Deserialize(dsr));
			dsr.ReadCollectionEndTag("Nodes");

			dsr.ReadEndTag("PModel");

			return new PIModel(inf, namemap, extmap, evalstack, nodes);
		}

		/// <summary>
		/// Compute a nice, slightly simplfied representation for display as DGML
		/// </summary>
		public String DGMLDumpPModel(int line)
		{
			StringBuilder xmlsb = new StringBuilder();
			System.Xml.XmlWriterSettings settings = new System.Xml.XmlWriterSettings { Indent = true, OmitXmlDeclaration = true, CloseOutput = true, Encoding = Encoding.UTF8 };
			using(var xw = XmlWriter.Create(xmlsb, settings))
			{
				xw.WriteStartDocument();
				xw.WriteStartElement("DirectedGraph", DgmlExtensions.NamespaceUri);

				xw.WriteStartElement("Nodes");
				
				for(int i = 0; i < this.Nodes.Count; ++i)
					this.Nodes[i].DGMLDumpHeapNode(line, xw);

				for(int i = 0; i < this.EvalStack.Count; ++i)
					this.EvalStack[i].DGMLDumpStackNode(i, xw);

				foreach(PIExternalRef extr in this.ExternalRefs.Keys)
				{
					xw.WriteStartElement("Node");
					xw.WriteAttributeString("Id", extr.ToString());
					xw.WriteAttributeString("Label", extr.ToString());
					xw.WriteEndElement();
				}

				xw.WriteEndElement();

				xw.WriteStartElement("Links");

				for(int i = 0; i < this.Nodes.Count; ++i)
				{
					PIMNode nn = this.Nodes[i];
					if(nn.OutEdges.Count != 0)
						PIMEdge.DGMLDumpHeapEdges(nn.OutEdges.ToArray(), nn, xw);
				}

				for(int i = 0; i < this.EvalStack.Count; ++i)
					this.EvalStack[i].DGMLDumpStackNodeEdges(i, xw);
				
				foreach(KeyValuePair<PIExternalRef, int[]> extr in this.ExternalRefs)
				{
					for(int i = 0; i < extr.Value.Length; ++i)
					{
						xw.WriteStartElement("Link");
						xw.WriteAttributeString("Id", extr.Key.ToString() + "_" + i.ToString());
						xw.WriteAttributeString("Source", extr.Key.ToString());
						xw.WriteAttributeString("Target", extr.Value[i].ToString());
						xw.WriteEndElement();
					}
				}

				xw.WriteEndElement();

				DgmlExtensions.WriteCategoriesAndProperties(xw);
				DgmlExtensions.WriteStyles(xw);
				
				xw.WriteEndElement();
				xw.WriteEndDocument();
			}

			return xmlsb.ToString();
		}

		/// <summary>
		/// Return the home node for the given name.
		/// </summary>
		public PIMNode ResolveHomeForRoot(IRRootIdentifier rrid)
		{ return this.Nodes[this.NameToNodeMap[rrid]]; }

		/// <summary>
		/// Return the eval stack node at the given position.
		/// </summary>
		public PISNode GetEvalStackNodeAtPos(int loc)
		{
			PISNode sn = this.EvalStack[this.EvalStack.Count - (loc + 1)];
			return sn;
		}

		/// <summary>
		/// Return a list of all the local variables (and live static roots) currently in scope.
		/// </summary>
		public List<IRRootIdentifier> GetAllLocalAndLiveStaticRoots()
		{ return new List<IRRootIdentifier>(this.NameToNodeMap.Keys); }

		/// <summary>
		/// Return a list of all the local variables currently in scope.
		/// </summary>
		public List<IRVarName> GetAllLocalVariables()
		{
			List<IRVarName> lvnl = new List<IRVarName>();
			
			foreach(IRRootIdentifier rrid in this.NameToNodeMap.Keys)
			{
				if(rrid is IRVarName)
					lvnl.Add((IRVarName)rrid);
			}

			return lvnl;
		}

		/// <summary>
		/// Return a list of all the static fields that may be read from or written to in this scope.
		/// </summary>
		public List<IRStaticField> GetAllLiveStaticFields()
		{
			List<IRStaticField> lvnl = new List<IRStaticField>();

			foreach(IRRootIdentifier rrid in this.NameToNodeMap.Keys)
			{
				if(rrid is IRStaticField)
					lvnl.Add((IRStaticField)rrid);
			}

			return lvnl;
		}

		/// <summary>
		/// Get a set of all the location identities for a node.
		/// </summary>
		public static HashSet<MemLocID> LocationIDSForNode(PIMNode n)
		{
			HashSet<MemLocID> mls = new HashSet<MemLocID>();
			AddAllLocationIDSForNodeToSet(n, mls);

			return mls;
		}

		/// <summary>
		/// Add all the location identities for a node to the given identity set.
		/// </summary>
		public static void AddAllLocationIDSForNodeToSet(PIMNode n, ICollection<MemLocID> identityset)
		{
			foreach(int id in n.Identities)
				identityset.Add(MemLocID.MakeExternalMemLocation(id));

			foreach(PIEscapeSite esci in n.EscapeLines)
				identityset.Add(MemLocID.MakeInternalMemLocation(esci.EscIdentity, esci.EscLine, n.HomeLocFor));
		}

		/// <summary>
		/// Extract the set of identity tags and escape sites for all the nodes that may be the target of the given variable/static name.
		/// </summary>
		public void GetIDSOfPossibleRootTargets(IRRootIdentifier rid, HashSet<MemLocID> identityset)
		{
			HashSet<PIMNode> mnl = this.GetPossibleRootTargetNodes(rid);
			foreach(PIMNode n in mnl)
				AddAllLocationIDSForNodeToSet(n, identityset);
		}

		/// <summary>
		/// Extract the set of nodes that may be the target of the given variable/static name.
		/// </summary>
		public HashSet<PIMNode> GetPossibleRootTargetNodes(IRRootIdentifier rid)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();
			PIMNode homenode = this.Nodes[this.NameToNodeMap[rid]];

			for(int i = 0; i < homenode.OutEdges.Count; ++i)
			{
				PIMNode trgt = this.Nodes[homenode.OutEdges[i].TargetAddress];
				mnl.Add(trgt);
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of identity tags and escape sites for all the nodes that may be reachable from the given variable/static name.
		/// </summary>
		public void GetIDSOfAllReachableNodesFromRoot(IRRootIdentifier rid, HashSet<MemLocID> identityset)
		{
			HashSet<PIMNode> mnl = this.GetAllReachableNodesFromRoot(rid);
			foreach(PIMNode n in mnl)
				AddAllLocationIDSForNodeToSet(n, identityset);
		}

		/// <summary>
		/// Extract the set of nodes that may be reachable from the given variable/static name.
		/// </summary>
		public HashSet<PIMNode> GetAllReachableNodesFromRoot(IRRootIdentifier rid)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();
			Queue<PIMNode> toproc = new Queue<PIMNode>();
			PIMNode homenode = this.Nodes[this.NameToNodeMap[rid]];

			for(int i = 0; i < homenode.OutEdges.Count; ++i)
			{
				PIMNode trgt = this.Nodes[homenode.OutEdges[i].TargetAddress];
				toproc.Enqueue(trgt); 
				mnl.Add(trgt);
			}

			while(toproc.Count != 0)
			{
				PIMNode nn = toproc.Dequeue();

				for(int i = 0; i < nn.OutEdges.Count; ++i)
				{
					PIMNode trgt = this.Nodes[nn.OutEdges[i].TargetAddress];

					if(!mnl.Contains(trgt))
						toproc.Enqueue(trgt);
					mnl.Add(trgt);
				}
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of nodes that may be reachable from the given eval stack location.
		/// </summary>
		public HashSet<PIMNode> GetPossibleEvalStackTargetNodes(int loc)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();

			PISNode sn = this.EvalStack[this.EvalStack.Count - (loc + 1)];
			for(int i = 0; i < sn.Targets.Length; ++i)
			{
				PIMNode trgt = this.Nodes[sn.Targets[i].TargetAddress];
				mnl.Add(trgt);
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of nodes that may be reachable from the given eval stack location.
		/// </summary>
		public HashSet<PIMNode> GetAllReachableNodesFromEvalStack(int loc)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();
			Queue<PIMNode> toproc = new Queue<PIMNode>();
			PISNode sn = this.EvalStack[this.EvalStack.Count - (loc + 1)];

			for(int i = 0; i < sn.Targets.Length; ++i)
			{
				PIMNode trgt = this.Nodes[sn.Targets[i].TargetAddress];
				toproc.Enqueue(trgt);
				mnl.Add(trgt);
			}

			while(toproc.Count != 0)
			{
				PIMNode nn = toproc.Dequeue();

				for(int i = 0; i < nn.OutEdges.Count; ++i)
				{
					PIMNode trgt = this.Nodes[nn.OutEdges[i].TargetAddress];

					if(!mnl.Contains(trgt))
						toproc.Enqueue(trgt);
					mnl.Add(trgt);
				}
			}

			return mnl;
		}

		/// <summary>
		/// Given a set of identity tags and escape sites return all the corresponding nodes.
		/// </summary>
		public HashSet<PIMNode> GetAllNodesAssociatedWithIDSet(HashSet<MemLocID> ids)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();

			for(int i = 0; i < this.Nodes.Count; ++i)
			{
				PIMNode n = this.Nodes[i];

				foreach(int id in n.Identities)
				{
					MemLocID lid = MemLocID.MakeExternalMemLocation(id);
					if(ids.Contains(lid))
						mnl.Add(n);
				}

				foreach(PIEscapeSite esci in n.EscapeLines)
				{
					MemLocID lid = MemLocID.MakeInternalMemLocation(esci.EscIdentity, esci.EscLine, n.HomeLocFor);
					if(ids.Contains(lid))
						mnl.Add(n);
				}
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of nodes that can be accessed via the load of the given field from any of the nodes in the input set.
		/// </summary>
		public HashSet<PIMNode> GetNodesFromSourceNodeAndField(HashSet<PIMNode> orignodes, IRMemberField f)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();

			foreach(PIMNode onode in orignodes)
			{
				for(int i = 0; i < onode.OutEdges.Count; ++i)
				{
					PIMEdge ee = onode.OutEdges[i];
					if(ee.OffsetInSource.PathLast == f)
					{
						PIMNode trgt = this.Nodes[ee.TargetAddress];
						mnl.Add(trgt);
					}
				}
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of nodes that may be reachable from the given node and field.
		/// </summary>
		public HashSet<PIMNode> GetAllReachableNodesSourceNodeAndField(PIMNode orignode, IRMemberField f)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>();
			Queue<PIMNode> toproc = new Queue<PIMNode>();

			for(int i = 0; i < orignode.OutEdges.Count; ++i)
			{
				PIMEdge ee = orignode.OutEdges[i];
				if(ee.OffsetInSource.PathLast == f)
				{
					PIMNode trgt = this.Nodes[ee.TargetAddress];
					toproc.Enqueue(trgt);
					mnl.Add(trgt);
				}
			}

			while(toproc.Count != 0)
			{
				PIMNode nn = toproc.Dequeue();

				for(int i = 0; i < nn.OutEdges.Count; ++i)
				{
					PIMNode trgt = this.Nodes[nn.OutEdges[i].TargetAddress];

					if(!mnl.Contains(trgt))
						toproc.Enqueue(trgt);
					mnl.Add(trgt);
				}
			}

			return mnl;
		}

		/// <summary>
		/// Extract the set of nodes that may be reachable from the given start set of nodes
		/// </summary>
		public HashSet<PIMNode> GetAllReachableNodesFromGivenStartSet(HashSet<PIMNode> startnodes)
		{
			HashSet<PIMNode> mnl = new HashSet<PIMNode>(startnodes);
			Queue<PIMNode> toproc = new Queue<PIMNode>(startnodes);
			
			while(toproc.Count != 0)
			{
				PIMNode nn = toproc.Dequeue();

				for(int i = 0; i < nn.OutEdges.Count; ++i)
				{
					PIMNode trgt = this.Nodes[nn.OutEdges[i].TargetAddress];

					if(!mnl.Contains(trgt))
						toproc.Enqueue(trgt);
					mnl.Add(trgt);
				}
			}

			return mnl;
		}
	}
}
